import vtk
import myVTKRenderWindowInteractorStyle
reload(myVTKRenderWindowInteractorStyle)


def InitializeVolume():
    """
    Setup a volume  for rendering.

    This has to be called before InitializeRegistration, since the DRR of
    the volume will be the "moving image" in the registration.
    """
    #Create the volume and volume reader
    #This might be a problem if we want to register a 3D object other than a volume (eg. polygonal mesh)
    filename = 'C:/Lib/jointtrack/branches/TestData/vert1.vtk'
    volume = vtk.vtkVolume()
    volumeReader = vtk.vtkStructuredPointsReader()
    volumeReader.SetFileName(filename)
    volumeReader.Update()

    #TODO:Allow the user to modify/customize the transfer functions to change how the DRR is rendered!!
    #Default settings for the opacity transfer function,
    opacityTransferFunction = vtk.vtkPiecewiseFunction()
    opacityTransferFunction.AddPoint(0.0, 0.0)
    opacityTransferFunction.AddPoint(70, 0.0)
    opacityTransferFunction.AddPoint(255, 0.5)

    # Create transfer mapping scalar value to color
    colorTransferFunction = vtk.vtkColorTransferFunction()
    colorTransferFunction.AddRGBPoint(0.0, 0.0, 0.0, 0.0)
    colorTransferFunction.AddRGBPoint(72, 1.0, 1.0, 1.0)
    colorTransferFunction.AddRGBPoint(255, 1.0, 1.0, 1.0)

    # The property describes how the data will look
    volumeProperty = vtk.vtkVolumeProperty()

    # Setup default properties
    volumeProperty.SetColor(colorTransferFunction)
    volumeProperty.SetScalarOpacity(opacityTransferFunction)
    volumeProperty.ShadeOn()
    volumeProperty.SetInterpolationTypeToLinear()

    volumeMapper = vtk.vtkVolumeTextureMapper3D()
    volumeMapper.SetInput(volumeReader.GetOutput())

    volume.SetProperty(volumeProperty)
    volume.SetMapper(volumeMapper)

    #Set up Moving Image 3D renderwindow, with interaction
    ren = vtk.vtkRenderer()
    #ren.AddVolume(volume)
    ren.AddActor(volume)
    rw = vtk.vtkRenderWindow()
    rw.AddRenderer(ren)
    iren = vtk.vtkRenderWindowInteractor()
    iren.SetRenderWindow(rw)
    rw.SetSize(512,512)

    style = myVTKRenderWindowInteractorStyle.myVTKRenderWindowInteractorStyle(iren)
    #style = vtk.vtkInteractorStyleTrackballCamera()
    iren.SetInteractorStyle(style)

    rw.Render()
    iren.Start()

if __name__ == "__main__":
    InitializeVolume()